import matplotlib.pyplot as plt
from scipy.interpolate import CubicSpline
import numpy as np
import pandas as pd

import rclpy
from rclpy.node import Node
from nav_msgs.msg import Path
from geometry_msgs.msg import PoseStamped
from tkinter import filedialog, messagebox
import tkinter as tk

"""
Use this node to create path:
    ros2 run path_creator path_creator
    select the start pose ,middle mark and target pose
    press 'f' to finish
    press 'r' to reset
    press 's' to save
    press 'z' to undo
    start rviz2 to see the path
"""


class PathCreator(Node):
    x: list[float]
    y: list[float]

    def __init__(self):
        super().__init__("path_creator")
        self.get_logger().info("Path creator started")

        self.declare_parameter("size", 100.0)
        self.size = self.get_parameter("size").value

        self.pathPub = self.create_publisher(Path, "reference_path", 3)
        self.create_timer(0.1, self.timer_callback)

        self.fig, self.ax = plt.subplots()
        self.x = [0]
        self.y = [0]
        plt.xlim(-self.size, self.size)
        plt.ylim(-self.size, self.size)
        self.ax.scatter(self.x, self.y)

        self.fig.canvas.mpl_disconnect(self.fig.canvas.manager.key_press_handler_id)
        self.fig.canvas.mpl_connect("button_press_event", self.on_button_press)
        self.fig.canvas.mpl_connect("key_press_event", self.on_key_press)
        self.fig.canvas.mpl_connect("close_event", self.on_close)

        plt.show(block=False)

    def timer_callback(self):
        plt.pause(0.1)

    def on_key_press(self, event):
        if event.key == "z":
            self.x.pop(-1)
            self.y.pop(-1)
            self.repaint()
        if event.key == "s":
            # save to file
            self.get_logger().info("Saving path to file")
            self.savePath()
        if event.key == "f":
            # finished
            if len(self.x) < 2 or len(self.y) < 2:
                return
            path_msg = Path()
            path_msg.header.frame_id = "map"
            path_msg.header.stamp = self.get_clock().now().to_msg()
            xs, ys = reInterpolate(self.x, self.y)
            for i in range(len(xs)):
                pose_stamped = PoseStamped()
                pose_stamped.header.frame_id = "map"
                pose_stamped.header.stamp = self.get_clock().now().to_msg()
                pose_stamped.pose.position.x = float(xs[i])
                pose_stamped.pose.position.y = float(ys[i])
                pose_stamped.pose.position.z = 0.0
                pose_stamped.pose.orientation.x = 0.0
                pose_stamped.pose.orientation.y = 0.0
                pose_stamped.pose.orientation.z = 0.0
                pose_stamped.pose.orientation.w = 1.0
                path_msg.poses.append(pose_stamped)
            self.pathPub.publish(path_msg)
        if event.key == "r":
            # reset
            self.x = [0]
            self.y = [0]
            self.repaint()

    def on_button_press(self, event):
        self.x.append(event.xdata)
        self.y.append(event.ydata)
        self.repaint()

    def repaint(self):
        self.ax.cla()
        plt.xlim(-self.size, self.size)
        plt.ylim(-self.size, self.size)
        self.ax.scatter(self.x, self.y)
        xs, ys = reInterpolate(self.x, self.y)
        self.ax.plot(xs, ys)
        for i in range(len(self.x)):
            self.ax.annotate(str(i), xy=(self.x[i], self.y[i]))
        self.fig.canvas.draw_idle()

    def savePath(self):
        root = tk.Tk()
        root.withdraw()
        interplot = messagebox.askyesno(
            "Save Mode", "Save raw path data? true=raw, false=interpolated"
        )
        path = filedialog.asksaveasfilename(
            defaultextension=".csv",
            filetypes=[("CSV files", "*.csv"), ("All files", "*.*")],
            title="Save Path",
        )
        if path:
            if interplot:
                data = pd.DataFrame({"x": self.x, "y": self.y})
            else:
                xs, ys = reInterpolate(self.x, self.y)
                data = pd.DataFrame({"x": xs, "y": ys})
            data.to_csv(path, index=False)
            self.get_logger().info("Path saved to " + path)
        else:
            self.get_logger().info("Path not saved")

    def on_close(self, event):
        self.get_logger().info("Path creator closed")
        rclpy.shutdown()


def reInterpolate(x: list, y: list) -> tuple[list, list]:
    if len(x) < 2 or len(y) < 2:
        return x, y

    new_x = np.arange(len(x))

    cs_x = CubicSpline(new_x, x)
    cs_y = CubicSpline(new_x, y)

    x_new = np.linspace(0, len(x) - 1, 100)
    x_interpolated = cs_x(x_new)
    y_interpolated = cs_y(x_new)

    return x_interpolated.tolist(), y_interpolated.tolist()


def main():
    rclpy.init()
    path_creator = PathCreator()
    rclpy.spin(path_creator)
    rclpy.shutdown()


if __name__ == "__main__":
    main()
